from utils import MainHandler, LLMType, EmbeddingType

# llama_index 文档参考：https://docs.llamaindex.ai/en/stable/


class CustomerServiceSystem:
    """
    电商客服系统
    """

    def __init__(self, llm_type: LLMType, embedding_type: EmbeddingType, **kwargs):
        self.main_handler = MainHandler(
            llm_type=llm_type,
            embedding_type=embedding_type,
            **kwargs,
        )

    def create_vector_database(self, knowledge_path: str):
        """
        创建向量数据库
        """
        self.main_handler.create_vector_database(knowledge_path)

    def create_query_engine(self, vector_database_path: str, streaming: bool = True):
        """
        创建查询引擎

        Args:
            vector_database_path: 向量数据库路径
            streaming: 是否使用流式输出，默认为True

        Returns:
            查询引擎实例
        """
        return self.main_handler.get_router_query_engine(
            vector_database_path, streaming=streaming
        )

    def create_vector_query_engine(
        self, vector_database_path: str, streaming: bool = True
    ):
        """
        创建普通向量查询引擎（不使用路由器）
        在路由器查询引擎失败时作为备用

        Args:
            vector_database_path: 向量数据库路径
            streaming: 是否使用流式输出，默认为True

        Returns:
            查询引擎实例
        """
        vector_tool = self.main_handler.get_vector_query_engine_tools(
            vector_database_path, streaming=streaming
        )
        return vector_tool.query_engine


if __name__ == "__main__":
    # 调用HuggingFace的本地LLM模型
    # mysystem = CustomerServiceSystem(
    #     LLMType.Local_HuggingFace,
    #     EmbeddingType.Local_HuggingFace,
    #     llm_model_name="deepseek-ai/DeepSeek-R1-Distill-Qwen-1.5B",
    #     device_map="cuda",
    # )

    # 调用OpenRouter的LLM模型
    # 调用免费模型的话每天有次数限制
    # mysystem = CustomerServiceSystem(
    #     LLMType.Online_OpenRouter,
    #     EmbeddingType.Local_HuggingFace,
    # )

    # 调用硅基流动的LLM模型
    mysystem = CustomerServiceSystem(
        LLMType.Online_OpenAILike,
        EmbeddingType.Local_HuggingFace,
        llm_model_name="deepseek-ai/DeepSeek-V3",
        api_base="https://api.siliconflow.cn/v1",
        is_function_calling_model=True,
    )

    # 首次创建向量数据库的时候运行，只需要运行一次
    # mysystem.create_vector_database("./knowledge")

    # 选择是否使用流式输出
    use_streaming = True  # 可以设置为False以使用非流式输出

    print(f"使用{'流式' if use_streaming else '非流式'}输出模式")

    # 创建查询引擎
    query_engine = mysystem.create_query_engine("./storage", streaming=use_streaming)

    # 创建备用向量查询引擎
    backup_query_engine = mysystem.create_vector_query_engine(
        "./storage", streaming=use_streaming
    )

    # 查询，添加异常处理
    query = "手机退货需要什么条件？"
    error_count = 0
    for i in range(1):
        try:
            # 尝试使用路由查询引擎
            response = query_engine.query(query)
            if use_streaming:
                print("\n查询回复结果：")
                for chuck in response.response_gen:
                    print(chuck, end="", flush=True)
            else:
                print(f"\n查询回复结果：{response.response}")

            for node in response.source_nodes:
                print(f"\n来源文档：{node.metadata.get('file_name')}")
                print(f"文档最后修改时间：{node.metadata.get('last_modified_date')}")
                print(f"置信度：{node.score:.2f}")

        except Exception as e:
            error_count += 1
            print(f"查询过程中发生错误: {type(e).__name__}: {e}")
            print("切换到备用向量查询引擎...")
            try:
                response = backup_query_engine.query(query)
                if use_streaming:
                    print("\n查询回复结果：")
                    for chuck in response.response_gen:
                        print(chuck, end="", flush=True)
                else:
                    print(f"\n查询回复结果：{response.response}")

                for node in response.source_nodes:
                    print(f"\n来源文档：{node.metadata.get('file_name')}")
                    print(
                        f"文档最后修改时间：{node.metadata.get('last_modified_date')}"
                    )
                    print(f"置信度：{node.score:.2f}")
            except Exception as e:
                print(f"备用引擎也失败了: {type(e).__name__}: {e}")

    print(f"错误次数：{error_count}")
